import sys
sys.path.append('../../')

from langchain.document_loaders import Docx2txtLoader, DirectoryLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import Chroma
from langchain.embeddings.huggingface import HuggingFaceEmbeddings
import sentence_transformers
from langchain.chains import RetrievalQA
from langchain.prompts import PromptTemplate

embedding_model_dict = {
    "ernie-tiny": "nghuyong/ernie-3.0-nano-zh",
    "ernie-base": "nghuyong/ernie-3.0-base-zh",
    "text2vec": "GanymedeNil/text2vec-large-chinese",
    "text2vec2":"uer/sbert-base-chinese-nli",
    "text2vec3":"shibing624/text2vec-base-chinese",
}
base_chroma = "../../chroma/nurse/"


EMBEDDING_MODEL = "ernie-tiny"
# 初始化 hugginFace 的 embeddings 对象
embeddings = HuggingFaceEmbeddings(model_name=embedding_model_dict[EMBEDDING_MODEL], )
embeddings.client = sentence_transformers.SentenceTransformer(
        embeddings.model_name, device='cpu')

db = Chroma(persist_directory='base_chroma', embedding_function=embeddings)

question = "三味书屋在哪里"

similarDocs = db.similarity_search(question, include_metadata=True,k=10)
for x in similarDocs:
    print(x)